/* Copyright 2020 Victor Penso

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>. */

package main

import (
        "io/ioutil"
        "os/exec"
        "log"
        "strings"
		"regexp"
        "strconv"
        "github.com/prometheus/client_golang/prometheus"
)

type PartitionNodesMetrics struct {
	alloc float64
	comp  float64
	down  float64
	drain float64
	err   float64
	fail  float64
	idle  float64
	maint float64
	mix   float64
	resv  float64
}

/*
func PartitionsData() []byte {
        cmd := exec.Command("sinfo", "-h", "-o%R,%C")
        stdout, err := cmd.StdoutPipe()
        if err != nil {
                log.Fatal(err)
        }
        if err := cmd.Start(); err != nil {
                log.Fatal(err)
        }
        out, _ := ioutil.ReadAll(stdout)
        if err := cmd.Wait(); err != nil {
                log.Fatal(err)
        }
        return out
}*/

// Execute the sinfo command and return its output
func PartitionNodesData() []byte {
	cmd := exec.Command("sinfo", "-h", "-o %D,%T,%R")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}
	out, _ := ioutil.ReadAll(stdout)
	if err := cmd.Wait(); err != nil {
		log.Fatal(err)
	}
	return out
}

//func PartitionNodesGetMetrics() *PartitionNodesMetrics {
//	return ParsePartitionNodesMetrics(PartitionNodesData())
//}
/*
func ParsePartitionsMetrics() map[string]*PartitionMetrics {
        partitions := make(map[string]*PartitionMetrics)
        lines := strings.Split(string(PartitionsData()), "\n")
        for _, line := range lines {
                if strings.Contains(line,",") {
                        // name of a partition
                        partition := strings.Split(line,",")[0]
                        _,key := partitions[partition]
                        if !key {
                                partitions[partition] = &PartitionMetrics{0,0,0,0,0,0}
                        }
                        states := strings.Split(line,",")[1]
                        allocated,_ := strconv.ParseFloat(strings.Split(states,"/")[0],64)
                        idle,_ := strconv.ParseFloat(strings.Split(states,"/")[1],64)
                        other,_ := strconv.ParseFloat(strings.Split(states,"/")[2],64)
                        total,_ := strconv.ParseFloat(strings.Split(states,"/")[3],64)
                        partitions[partition].allocated = allocated
                        partitions[partition].idle = idle
                        partitions[partition].other = other
                        partitions[partition].total = total
                }
        }
        // get list of pending jobs by partition name
        list := strings.Split(string(PartitionsPendingJobsData()),"\n")
        for _,partition := range list {
		// accumulate the number of pending jobs
		_,key := partitions[partition]
		if key {
				partitions[partition].pending += 1
            }
        }
		
		rlist := strings.Split(string(PartitionsRunningJobsData()),"\n")
        for _,partition := range rlist {
		// accumulate the number of running jobs
		_,key := partitions[partition]
		if key {
			partitions[partition].running += 1
            }
        }
        return partitions
}*/

//map[string]*PartitionNodesMetrics：定义map键值对类型，键为string数组，值为PartitionNodesMetrics指针

func ParsePartitionNodesMetrics() map[string]*PartitionNodesMetrics{
	//make，创建空的键值对类型
	partitions := make(map[string]*PartitionNodesMetrics)
	//字符串分割
	lines := strings.Split(string(PartitionNodesData()), "\n")
	// range: 循环遍历line
	for _, line := range lines {
		if strings.Contains(line,",") {
			split := strings.Split(line,",");
			partition := split[2]
			_,key := partitions[partition]
			//判断key是否存在，如果不存在，则分配map内存
			if !key {                    
                partitions[partition] = &PartitionNodesMetrics{0,0,0,0,0,0,0,0,0,0}
				count, _ := strconv.ParseFloat(strings.TrimSpace(split[0]), 64)
				state := split[1]
				// 创建正则匹配表达式
				alloc := regexp.MustCompile(`^alloc`)
				comp := regexp.MustCompile(`^comp`)
				down := regexp.MustCompile(`^down`)
				drain := regexp.MustCompile(`^drain`)
				fail := regexp.MustCompile(`^fail`)
				err := regexp.MustCompile(`^err`)
				idle := regexp.MustCompile(`^idle`)
				maint := regexp.MustCompile(`^maint`)
				mix := regexp.MustCompile(`^mix`)
				resv := regexp.MustCompile(`^res`)
				switch {
				case alloc.MatchString(state) == true:
					partitions[partition].alloc = count
				case comp.MatchString(state) == true:
					partitions[partition].comp = count
				case down.MatchString(state) == true:
					partitions[partition].down = count
				case drain.MatchString(state) == true:
					partitions[partition].drain = count
				case fail.MatchString(state) == true:
					partitions[partition].fail = count
				case err.MatchString(state) == true:
					partitions[partition].err = count
				case idle.MatchString(state) == true:
					partitions[partition].idle = count
				case maint.MatchString(state) == true:
					partitions[partition].maint = count
				case mix.MatchString(state) == true:
					partitions[partition].mix = count
				case resv.MatchString(state) == true:
					partitions[partition].resv = count
				}
            }else{
				count, _ := strconv.ParseFloat(strings.TrimSpace(split[0]), 64)
				state := split[1]
				alloc := regexp.MustCompile(`^alloc`)
				comp := regexp.MustCompile(`^comp`)
				down := regexp.MustCompile(`^down`)
				drain := regexp.MustCompile(`^drain`)
				fail := regexp.MustCompile(`^fail`)
				err := regexp.MustCompile(`^err`)
				idle := regexp.MustCompile(`^idle`)
				maint := regexp.MustCompile(`^maint`)
				mix := regexp.MustCompile(`^mix`)
				resv := regexp.MustCompile(`^res`)
				switch {
				case alloc.MatchString(state) == true:
					partitions[partition].alloc = count
				case comp.MatchString(state) == true:
					partitions[partition].comp = count
				case down.MatchString(state) == true:
					partitions[partition].down = count
				case drain.MatchString(state) == true:
					partitions[partition].drain = count
				case fail.MatchString(state) == true:
					partitions[partition].fail = count
				case err.MatchString(state) == true:
					partitions[partition].err = count
				case idle.MatchString(state) == true:
					partitions[partition].idle = count
				case maint.MatchString(state) == true:
					partitions[partition].maint = count
				case mix.MatchString(state) == true:
					partitions[partition].mix = count
				case resv.MatchString(state) == true:
					partitions[partition].resv = count
				}
			}
		}
	}
	return partitions
}


type PartitionNodesCollector struct {
	alloc *prometheus.Desc
	comp  *prometheus.Desc
	down  *prometheus.Desc
	drain *prometheus.Desc
	err   *prometheus.Desc
	fail  *prometheus.Desc
	idle  *prometheus.Desc
	maint *prometheus.Desc
	mix   *prometheus.Desc
	resv  *prometheus.Desc
}

func NewPartitionNodesCollector() *PartitionNodesCollector {
	labels := []string{"partition"}
	return &PartitionNodesCollector{
		alloc: prometheus.NewDesc("slurm_partition_nodes_alloc", "Allocated nodes", labels, nil),
		comp:  prometheus.NewDesc("slurm_partition_nodes_comp", "Completing nodes", labels, nil),
		down:  prometheus.NewDesc("slurm_partition_nodes_down", "Down nodes", labels, nil),
		drain: prometheus.NewDesc("slurm_partition_nodes_drain", "Drain nodes", labels, nil),
		err:   prometheus.NewDesc("slurm_partition_nodes_err", "Error nodes", labels, nil),
		fail:  prometheus.NewDesc("slurm_partition_nodes_fail", "Fail nodes", labels, nil),
		idle:  prometheus.NewDesc("slurm_partition_nodes_idle", "Idle nodes", labels, nil),
		maint: prometheus.NewDesc("slurm_partition_nodes_maint", "Maint nodes", labels, nil),
		mix:   prometheus.NewDesc("slurm_partition_nodes_mix", "Mix nodes", labels, nil),
		resv:  prometheus.NewDesc("slurm_partition_nodes_resv", "Reserved nodes", labels, nil),
	}
}

func (pnc *PartitionNodesCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- pnc.alloc
	ch <- pnc.comp
	ch <- pnc.down
	ch <- pnc.drain
	ch <- pnc.err
	ch <- pnc.fail
	ch <- pnc.idle
	ch <- pnc.maint
	ch <- pnc.mix
	ch <- pnc.resv
}

func (pnc *PartitionNodesCollector) Collect(ch chan<- prometheus.Metric) {
        pm := ParsePartitionNodesMetrics()
        for p := range pm {
                if pm[p].alloc > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.alloc, prometheus.GaugeValue, pm[p].alloc, p)
                }
                if pm[p].comp > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.comp, prometheus.GaugeValue, pm[p].comp, p)
                }
                if pm[p].down > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.down, prometheus.GaugeValue, pm[p].down, p)
                }
                if pm[p].drain > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.drain, prometheus.GaugeValue, pm[p].drain, p)
                }
				if pm[p].err > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.err, prometheus.GaugeValue, pm[p].err, p)
                }
                if pm[p].fail > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.fail, prometheus.GaugeValue, pm[p].fail, p)
                }
				if pm[p].idle > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.idle, prometheus.GaugeValue, pm[p].idle, p)
                }
                if pm[p].maint > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.maint, prometheus.GaugeValue, pm[p].maint, p)
                }
				if pm[p].mix > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.mix, prometheus.GaugeValue, pm[p].mix, p)
                }
                if pm[p].resv > 0 {
                        ch <- prometheus.MustNewConstMetric(pnc.resv, prometheus.GaugeValue, pm[p].resv, p)
                }
        }
}
